package org.test.rocketmq.producer05;

import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 事务消息发送。在普通消息基础上，支持二阶段的提交能力。将二阶段提交和本地事务绑定，实现全局提交结果的一致性。
 *
 * <p>
 *     transactionCheckMax: 事务最大重试次数，如果超过检测次数，消息会默认为丢弃，即回滚消息
 *     transactionCheckInterval: 每隔多少毫秒进行重试检查，默认 60000(ms)
 *     transactionTimeOut: 6 min 判断消息超时进入回查的时间
 * </p>
 */
public class TransactionProducerExample {

    public static void main(String[] args) throws MQClientException, InterruptedException {
        TransactionListener transactionListener = new TransactionListenerImpl();
        TransactionMQProducer producer = new TransactionMQProducer("transaction_producer_unique_group_name");
        producer.setNamesrvAddr("localhost:9876");
        ExecutorService executorService = new ThreadPoolExecutor(2, 5, 100, TimeUnit.SECONDS, new ArrayBlockingQueue<>(2000),
                r -> {
                    Thread thread = new Thread(r);
                    thread.setName("client-transaction-msg-check-thread");
                    return thread;
                });
        producer.setExecutorService(executorService);
        producer.setTransactionListener(transactionListener);
        producer.start();


        String[] tags = new String[]{"TagA", "TagB", "TagC", "TagD", "TagE"};
        for (int i = 0; i < 10; i++) {
            Message msg =
                    new Message("TransactionTopicTest", tags[i % tags.length], "KeyTrans" + i,
                            ("Hello RocketMQ " + i).getBytes(StandardCharsets.UTF_8));
            SendResult sendResult = producer.sendMessageInTransaction(msg, null);
            System.out.printf("%s%n", sendResult);

            Thread.sleep(10);
        }
        for (int i = 0; i < 100000; i++) {
            Thread.sleep(1000);
        }
        producer.shutdown();
    }


    static class TransactionListenerImpl implements TransactionListener {

        private final AtomicInteger transactionIndex = new AtomicInteger(0);

        private final ConcurrentHashMap<String, Integer> localTrans = new ConcurrentHashMap<>();

        @Override
        public LocalTransactionState executeLocalTransaction(Message msg, Object arg) {
            System.out.printf("执行本地事务 %s %n", msg.getTransactionId());
            int value = transactionIndex.getAndIncrement();
            int status = value % 3;
            localTrans.put(msg.getTransactionId(), status);
            return LocalTransactionState.UNKNOW;
        }

        @Override
        public LocalTransactionState checkLocalTransaction(MessageExt msg) {
            System.out.printf("检查本地事务 %s %n", msg.getTransactionId());
            Integer status = localTrans.get(msg.getTransactionId());
            if (null != status) {
                return switch (status) {
                    // 暂时无法判断状态，等待固定时间以后Broker端根据回查规则向生产者进行消息回查
                    case 0 -> LocalTransactionState.UNKNOW;
                    // 提交事务，允许消费者消费该消息
                    case 1 -> LocalTransactionState.COMMIT_MESSAGE;
                    // 回滚事务，消息将被丢弃不允许消费
                    case 2 -> LocalTransactionState.ROLLBACK_MESSAGE;
                    default -> LocalTransactionState.COMMIT_MESSAGE;
                };
            }
            return LocalTransactionState.COMMIT_MESSAGE;
        }
    }
}
